/*M///////////////////////////////////////////////////////////////////////////////////////
//
//  IMPORTANT: READ BEFORE DOWNLOADING, COPYING, INSTALLING OR USING.
//
//  By downloading, copying, installing or using the software you agree to this license.
//  If you do not agree to this license, do not download, install,
//  copy or use the software.
//
//
//                        Intel License Agreement
//
// Copyright (C) 2000, Intel Corporation, all rights reserved.
// Third party copyrights are property of their respective owners.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
//   * Redistribution's of source code must retain the above copyright notice,
//     this list of conditions and the following disclaimer.
//
//   * Redistribution's in binary form must reproduce the above copyright notice,
//     this list of conditions and the following disclaimer in the documentation
//     and/or other materials provided with the distribution.
//
//   * The name of Intel Corporation may not be used to endorse or promote products
//     derived from this software without specific prior written permission.
//
// This software is provided by the copyright holders and contributors "as is" and
// any express or implied warranties, including, but not limited to, the implied
// warranties of merchantability and fitness for a particular purpose are disclaimed.
// In no event shall the Intel Corporation or contributors be liable for any direct,
// indirect, incidental, special, exemplary, or consequential damages
// (including, but not limited to, procurement of substitute goods or services;
// loss of use, data, or profits; or business interruption) however caused
// and on any theory of liability, whether in contract, strict liability,
// or tort (including negligence or otherwise) arising in any way out of
// the use of this software, even if advised of the possibility of such damage.
//
//M*/

#include "precomp.hpp"

static float CalcAverageMask(CvBlob* pBlob, IplImage* pImgFG ) {
	/* Calculate sum of mask: */
	double  Area, Aver = 0;
	CvRect  r;
	CvMat   mat;

	if (pImgFG == NULL) { return 0; }

	r.x = cvRound(pBlob->x - pBlob->w * 0.5);
	r.y = cvRound(pBlob->y - pBlob->h * 0.5);
	r.width = cvRound(pBlob->w);
	r.height = cvRound(pBlob->h);
	Area = r.width * r.height;
	if (r.x < 0) {r.width += r.x; r.x = 0;}
	if (r.y < 0) {r.height += r.y; r.y = 0;}
	if ((r.x + r.width) >= pImgFG->width) {r.width = pImgFG->width - r.x - 1;}
	if ((r.y + r.height) >= pImgFG->height) {r.height = pImgFG->height - r.y - 1;}

	if (r.width > 0 && r.height > 0) {
		double Sum = cvSum(cvGetSubRect(pImgFG, &mat, r)).val[0] / 255.0;
		assert(Area > 0);
		Aver = Sum / Area;
	}
	return (float)Aver;
}   /* Calculate sum of mask. */


/*============== BLOB TRACKERCC CLASS DECLARATION =============== */
typedef struct DefBlobTracker {
	CvBlob                      blob;
	CvBlobTrackPredictor*       pPredictor;
	CvBlob                      BlobPredict;
	int                         Collision;
	CvBlobSeq*                  pBlobHyp;
	float                       AverFG;
} DefBlobTracker;

void cvFindBlobsByCCClasters(IplImage* pFG, CvBlobSeq* pBlobs, CvMemStorage* storage);

class CvBlobTrackerCC : public CvBlobTracker {
private:
	float           m_AlphaSize;
	float           m_AlphaPos;
	float           m_Alpha;
	int             m_Collision;
	int             m_ConfidenceType;
	const char*           m_ConfidenceTypeStr;
	CvBlobSeq       m_BlobList;
	CvBlobSeq       m_BlobListNew;
//  int             m_LastID;
	CvMemStorage*   m_pMem;
	int             m_ClearHyp;
	IplImage*       m_pImg;
	IplImage*       m_pImgFG;
public:
	CvBlobTrackerCC(): m_BlobList(sizeof(DefBlobTracker)) {
//      m_LastID = 0;
		m_ClearHyp = 0;
		m_pMem = cvCreateMemStorage();
		m_Collision = 1; /* if 1 then collistion will be detected and processed */
		AddParam("Collision", &m_Collision);
		CommentParam("Collision", "If 1 then collision cases are processed in special way");

		m_AlphaSize = 0.02f;
		AddParam("AlphaSize", &m_AlphaSize);
		CommentParam("AlphaSize", "Size update speed (0..1)");

		m_AlphaPos = 1.0f;
		AddParam("AlphaPos", &m_AlphaPos);
		CommentParam("AlphaPos", "Position update speed (0..1)");

		m_Alpha = 0.001f;
		AddParam("Alpha", &m_Alpha);
		CommentParam("Alpha", "Coefficient for model histogram updating (0 - hist is not updated)");

		m_ConfidenceType = 0;
		m_ConfidenceTypeStr = "NearestBlob";
		AddParam("ConfidenceType", &m_ConfidenceTypeStr);
		CommentParam("ConfidenceType", "Type of calculated Confidence (NearestBlob, AverFG, BC)");

		SetModuleName("CC");
	};

	~CvBlobTrackerCC() {
		if (m_pMem) { cvReleaseMemStorage(&m_pMem); }
	};

	/* Blob functions: */
	virtual int     GetBlobNum() {return m_BlobList.GetBlobNum();};
	virtual CvBlob* GetBlob(int BlobIndex) {return m_BlobList.GetBlob(BlobIndex);};
	virtual void    SetBlob(int BlobIndex, CvBlob* pBlob) {
		CvBlob* pB = m_BlobList.GetBlob(BlobIndex);
		if (pB) { pB[0] = pBlob[0]; }
	};

	virtual CvBlob* GetBlobByID(int BlobID) {return m_BlobList.GetBlobByID(BlobID);};
	virtual void    DelBlob(int BlobIndex) {
		DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(BlobIndex);
		if (pBT == NULL) { return; }
		if (pBT->pPredictor) {
			pBT->pPredictor->Release();
		} else {
			printf("WARNING!!! Invalid Predictor in CC tracker");
		}
		delete pBT->pBlobHyp;
		m_BlobList.DelBlob(BlobIndex);
	};
#if 0
	virtual void    DelBlobByID(int BlobID) {
		DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlobByID(BlobID);
		pBT->pPredictor->Release();
		delete pBT->pBlobHyp;
		m_BlobList.DelBlobByID(BlobID);
	};
#endif
	virtual void    Release() {delete this;};

	/* Add new blob to track it and assign to this blob personal ID */
	/* pBlob - pinter to structure with blob parameters (ID is ignored)*/
	/* pImg - current image */
	/* pImgFG - current foreground mask */
	/* return pointer to new added blob */
	virtual CvBlob* AddBlob(CvBlob* pB, IplImage* /*pImg*/, IplImage* pImgFG = NULL ) {
		assert(pImgFG); /* This tracker uses only foreground mask. */
		DefBlobTracker NewB;
		NewB.blob = pB[0];
//        CV_BLOB_ID(&NewB) = m_LastID;
		NewB.pBlobHyp = new CvBlobSeq;
		NewB.pPredictor = cvCreateModuleBlobTrackPredictKalman(); /* Module for position prediction. */
		NewB.pPredictor->Update(pB);
		NewB.AverFG = pImgFG ? CalcAverageMask(pB, pImgFG) : 0;
		m_BlobList.AddBlob((CvBlob*)&NewB);
		return m_BlobList.GetBlob(m_BlobList.GetBlobNum() - 1);
	};

	virtual void    Process(IplImage* pImg, IplImage* pImgFG = NULL) {
		CvSeq*      cnts;
		CvSeq*      cnt;
		int i;

		m_pImg = pImg;
		m_pImgFG = pImgFG;

		if (m_BlobList.GetBlobNum() <= 0 ) { return; }

		/* Clear bloblist for new blobs: */
		m_BlobListNew.Clear();

		assert(m_pMem);
		cvClearMemStorage(m_pMem);
		assert(pImgFG);


		/* Find CC: */
#if 0
		{
			// By contour clustering:
			cvFindBlobsByCCClasters(pImgFG, &m_BlobListNew, m_pMem);
		}
#else
		{
			/* One contour - one blob: */
			IplImage* pBin = cvCloneImage(pImgFG);
			assert(pBin);
			cvThreshold(pBin, pBin, 128, 255, CV_THRESH_BINARY);
			cvFindContours(pBin, m_pMem, &cnts, sizeof(CvContour), CV_RETR_EXTERNAL);

			/* Process each contour: */
			for (cnt = cnts; cnt; cnt = cnt->h_next) {
				CvBlob  NewBlob;

				/* Image moments: */
				double      M00, X, Y, XX, YY;
				CvMoments   m;
				CvRect      r = ((CvContour*)cnt)->rect;
				CvMat       mat;
				if (r.height < 3 || r.width < 3) { continue; }
				cvMoments( cvGetSubRect(pImgFG, &mat, r), &m, 0 );
				M00 = cvGetSpatialMoment( &m, 0, 0 );
				if (M00 <= 0 ) { continue; }
				X = cvGetSpatialMoment( &m, 1, 0 ) / M00;
				Y = cvGetSpatialMoment( &m, 0, 1 ) / M00;
				XX = (cvGetSpatialMoment( &m, 2, 0 ) / M00) - X * X;
				YY = (cvGetSpatialMoment( &m, 0, 2 ) / M00) - Y * Y;
				NewBlob = cvBlob(r.x + (float)X, r.y + (float)Y, (float)(4 * sqrt(XX)), (float)(4 * sqrt(YY)));
				m_BlobListNew.AddBlob(&NewBlob);
			}   /* Next contour. */

			cvReleaseImage(&pBin);
		}
#endif
		for (i = m_BlobList.GetBlobNum(); i > 0; --i) {
			/* Predict new blob position: */
			CvBlob*         pB = NULL;
			DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(i - 1);

			/* Update predictor by previous value of blob: */
			pBT->pPredictor->Update(&(pBT->blob));

			/* Predict current position: */
			pB = pBT->pPredictor->Predict();

			if (pB) {
				pBT->BlobPredict = pB[0];
			} else {
				pBT->BlobPredict = pBT->blob;
			}
		}   /* Predict new blob position. */

		if (m_Collision)
			for (i = m_BlobList.GetBlobNum(); i > 0; --i) {
				/* Predict collision. */
				int             Collision = 0;
				int             j;
				DefBlobTracker* pF = (DefBlobTracker*)m_BlobList.GetBlob(i - 1);

				for (j = m_BlobList.GetBlobNum(); j > 0; --j) {
					/* Predict collision: */
					CvBlob* pB1;
					CvBlob* pB2;
					DefBlobTracker* pF2 = (DefBlobTracker*)m_BlobList.GetBlob(j - 1);
					if (i == j) { continue; }
					pB1 = &pF->BlobPredict;
					pB2 = &pF2->BlobPredict;

					if ( fabs(pB1->x - pB2->x) < 0.6 * (pB1->w + pB2->w) &&
							fabs(pB1->y - pB2->y) < 0.6 * (pB1->h + pB2->h) ) { Collision = 1; }

					pB1 = &pF->blob;
					pB2 = &pF2->blob;

					if ( fabs(pB1->x - pB2->x) < 0.6 * (pB1->w + pB2->w) &&
							fabs(pB1->y - pB2->y) < 0.6 * (pB1->h + pB2->h) ) { Collision = 1; }

					if (Collision) { break; }

				}   /* Check next blob to cross current. */

				pF->Collision = Collision;

			}   /* Predict collision. */

		for (i = m_BlobList.GetBlobNum(); i > 0; --i) {
			/* Find a neighbour on current frame
			 * for each blob from previous frame:
			 */
			CvBlob*         pB = m_BlobList.GetBlob(i - 1);
			DefBlobTracker* pBT = (DefBlobTracker*)pB;
			//int             BlobID = CV_BLOB_ID(pB);
			//CvBlob*         pBBest = NULL;
			//double          DistBest = -1;
			//int j;

			if (pBT->pBlobHyp->GetBlobNum() > 0) {
				/* Track all hypotheses: */
				int h, hN = pBT->pBlobHyp->GetBlobNum();
				for (h = 0; h < hN; ++h) {
					int         j, jN = m_BlobListNew.GetBlobNum();
					CvBlob*     pB = pBT->pBlobHyp->GetBlob(h);
					int         BlobID = CV_BLOB_ID(pB);
					CvBlob*     pBBest = NULL;
					double      DistBest = -1;
					for (j = 0; j < jN; j++) {
						/* Find best CC: */
						double  Dist = -1;
						CvBlob* pBNew = m_BlobListNew.GetBlob(j);
						double  dx = fabs(CV_BLOB_X(pB) - CV_BLOB_X(pBNew));
						double  dy = fabs(CV_BLOB_Y(pB) - CV_BLOB_Y(pBNew));
						if (dx > 2 * CV_BLOB_WX(pB) || dy > 2 * CV_BLOB_WY(pB)) { continue; }

						Dist = sqrt(dx * dx + dy * dy);
						if (Dist < DistBest || pBBest == NULL) {
							DistBest = Dist;
							pBBest = pBNew;
						}
					}   /* Find best CC. */

					if (pBBest) {
						pB[0] = pBBest[0];
						CV_BLOB_ID(pB) = BlobID;
					} else {
						/* Delete this hypothesis. */
						pBT->pBlobHyp->DelBlob(h);
						h--;
						hN--;
					}
				}   /* Next hypothysis. */
			}   /*  Track all hypotheses. */
		}   /*  Track next blob. */

		m_ClearHyp = 1;

	} /* Process. */

	virtual void ProcessBlob(int BlobIndex, CvBlob* pBlob, IplImage* /*pImg*/, IplImage* /*pImgFG*/ = NULL) {
		int             ID = pBlob->ID;
		CvBlob*         pB = m_BlobList.GetBlob(BlobIndex);
		DefBlobTracker* pBT = (DefBlobTracker*)pB;
		//CvBlob*         pBBest = NULL;
		//double          DistBest = -1;
		int             BlobID;

		if (pB == NULL) { return; }

		BlobID = pB->ID;

		if (m_Collision && pBT->Collision) {
			/* Tracking in collision: */
			pB[0] = pBT->BlobPredict;
			CV_BLOB_ID(pB) = BlobID;
		}   /* Tracking in collision. */
		else {
			/* Non-collision tracking: */
			CvBlob* pBBest = GetNearestBlob(pB);
			if (pBBest) {
				float   w = pBlob->w * (1 - m_AlphaSize) + m_AlphaSize * pBBest->w;
				float   h = pBlob->h * (1 - m_AlphaSize) + m_AlphaSize * pBBest->h;
				float   x = pBlob->x * (1 - m_AlphaPos) + m_AlphaPos * pBBest->x;
				float   y = pBlob->y * (1 - m_AlphaPos) + m_AlphaPos * pBBest->y;
				pB->w = w;
				pB->h = h;
				pB->x = x;
				pB->y = y;
				CV_BLOB_ID(pB) = BlobID;
			}
		}   /* Non-collision tracking. */

		pBlob[0] = pB[0];
		pBlob->ID = ID;
	};

	virtual double  GetConfidence(int BlobIndex, CvBlob* pBlob, IplImage* /*pImg*/, IplImage* pImgFG = NULL) {
		/* Define coefficients in exp by exp(-XT*K)=VT: */
		static double _KS = -log(0.1) / pow(0.5, 2); /* XT = 1, VT = 0.1 - when size is Larger in 2 times Confidence is smoller in 10 times */
		static double _KP = -log(0.1) / pow(m_pImg->width * 0.02, 2); /* XT = 0.02*ImgWidth, VT = 0.1*/
		DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(BlobIndex);
		float   dx, dy, dw, dh;
		float   dp2, ds2;
		double  W = 1;
		CvBlob* pBC = GetNearestBlob(pBlob);
		if (pBC == NULL ) { return 0; }

		dx = pBC->x - pBlob->x;
		dy = pBC->y - pBlob->y;
		dw = (pBC->w - pBlob->w) / pBC->w;
		dh = (pBC->h - pBlob->h) / pBC->h;

		dp2 = dx * dx + dy * dy;
		ds2 = dw * dw + dh * dh;

		if (!pBT->Collision) {
			/* Confidence for size by nearest blob: */
			W *= exp(-_KS * ds2);
		}

		if (m_ConfidenceType == 0 && !pBT->Collision) {
			/* Confidence by nearest blob: */
			W *= exp(-_KP * dp2);
		}

		if (m_ConfidenceType == 1 && pBT->AverFG > 0) {
			/* Calculate sum of mask: */
			float   Aver = CalcAverageMask(pBlob, pImgFG );
			if (Aver < pBT->AverFG) {
				float diff = 1 + 0.9f * (Aver - pBT->AverFG) / pBT->AverFG;
				if (diff < 0.1f) { diff = 0.1f; }
				W *= diff;
			}
		}   /* Calculate sum of mask. */

		if (m_ConfidenceType == 2) {
			/* Calculate BCoeff: */
			float   S = 0.2f;
			float   Aver = CalcAverageMask(pBlob, pImgFG );
			double B = sqrt(Aver * pBT->AverFG) + sqrt((1 - Aver) * (1 - pBT->AverFG));

			W *= exp((B - 1) / (2 * S));
		}   /* Calculate sum of mask. */

		return W;
	};

	virtual void UpdateBlob(int BlobIndex, CvBlob* /*pBlob*/, IplImage* /*pImg*/, IplImage* pImgFG = NULL) {
		DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(BlobIndex);

		if (pImgFG == NULL || pBT == NULL) { return; }

		if (!pBT->Collision) {
			//pBT->AverFG = pBT->AverFG * (1-m_Alpha) + m_Alpha * CalcAverageMask(pBlob,pImgFG);
		}
	};

	virtual void ParamUpdate() {
		const char*   pCT[3] = {"NearestBlob", "AverFG", "BC"};
		int     i;

		CvBlobTracker::ParamUpdate();

		for (i = 0; i < 3; ++i) {
			if (cv_stricmp(m_ConfidenceTypeStr, pCT[i]) == 0) {
				m_ConfidenceType = i;
			}
		}
		SetParamStr("ConfidenceType", pCT[m_ConfidenceType]);
	}

	/* ===============  MULTI HYPOTHESIS INTERFACE ================== */
	/* Return number of position hypotheses of currently tracked blob: */
	virtual int     GetBlobHypNum(int BlobIdx) {
		DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(BlobIdx);
		assert(pBT->pBlobHyp);
		return pBT->pBlobHyp->GetBlobNum();
	};  /* CvBlobtrackerList::GetBlobHypNum() */

	/* Return pointer to specified blob hypothesis by index blob: */
	virtual CvBlob* GetBlobHyp(int BlobIndex, int hypothesis) {
		DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(BlobIndex);
		assert(pBT->pBlobHyp);
		return pBT->pBlobHyp->GetBlob(hypothesis);
	};  /* CvBlobtrackerList::GetBlobHyp() */

	/* Set new parameters for specified (by index) blob hypothesis
	 * (can be called several times for each hypothesis):
	 */
	virtual void    SetBlobHyp(int BlobIndex, CvBlob* pBlob) {
		if (m_ClearHyp) {
			/* Clear all hypotheses: */
			int b, bN = m_BlobList.GetBlobNum();
			for (b = 0; b < bN; ++b) {
				DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(b);
				assert(pBT->pBlobHyp);
				pBT->pBlobHyp->Clear();
			}
			m_ClearHyp = 0;
		}
		{
			/* Add hypothesis: */
			DefBlobTracker* pBT = (DefBlobTracker*)m_BlobList.GetBlob(BlobIndex);
			assert(pBT->pBlobHyp);
			pBT->pBlobHyp->AddBlob(pBlob);
		}
	};

private:
	CvBlob* GetNearestBlob(CvBlob* pB) {
		//DefBlobTracker* pBT = (DefBlobTracker*)pB;
		CvBlob*         pBBest = NULL;
		double          DistBest = -1;
		int             j, BlobID;

		if (pB == NULL) { return NULL; }

		BlobID = pB->ID;

		for (j = m_BlobListNew.GetBlobNum(); j > 0; --j) {
			/* Find best CC: */
			double  Dist = -1;
			CvBlob* pBNew = m_BlobListNew.GetBlob(j - 1);
			double  dx = fabs(CV_BLOB_X(pB) - CV_BLOB_X(pBNew));
			double  dy = fabs(CV_BLOB_Y(pB) - CV_BLOB_Y(pBNew));
			if (dx > 2 * CV_BLOB_WX(pB) || dy > 2 * CV_BLOB_WY(pB)) { continue; }

			Dist = sqrt(dx * dx + dy * dy);
			if (Dist < DistBest || pBBest == NULL) {
				DistBest = Dist;
				pBBest = pBNew;
			}

		}   /* Find best CC. */

		return pBBest;

	}; /* GetNearestBlob */

};

CvBlobTracker* cvCreateBlobTrackerCC() {
	return (CvBlobTracker*) new CvBlobTrackerCC;
}
/*============== BLOB TRACKERCC CLASS DECLARATION =============== */
